# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

import pandas as pd
from data_formulator.agents.agent_utils import generate_data_summary, extract_code_from_gpt_response

import logging

logger = logging.getLogger(__name__)


SYSTEM_PROMPT = '''您是一位帮助用户解释代码的数据科学家,
需要让非技术人员能清晰理解代码功能。您将获得输入数据的摘要和转换代码。
请生成分步的代码解释, 让读者从逻辑层面理解数据转换的高阶步骤。
请保持简洁并维持高阶视角。读者不需要理解代码细节, 只想了解大致的处理逻辑。

重点解释新字段的计算方式, 不要解释 "return"、"load data" 等基础操作。

使用 Markdown 格式进行解释: 高亮常量、数据字段和重要动词。务必保持简洁。


For example:

[CONTEXT]

数据集、字段摘要及样例数据:

table_0 (us_covid_cases) fields:
	 Date -- type: object, values: 1/1/2021, 1/1/2022, 1/1/2023, ..., 9/8/2022, 9/9/2020, 9/9/2021, 9/9/2022
	Cases -- type:  int64, values:   -23999,   -14195,    -6940, ...,  1018935,  1032159,  1178403,  1433977

table_0 (us_covid_cases) sample:
```
|Date|Cases
0|1/21/2020|1
1|1/22/2020|0
2|1/23/2020|0
3|1/24/2020|1
4|1/25/2020|1
......
```

[CODE]

```python
import pandas as pd

def transform_data(df_0):
    # 转换: 将 Date 字段转为 datetime 格式
    df_0['Date'] = pd.to_datetime(df_0['Date'])

    # 排序: 按 Date 列排序 DataFrame 数据
    df_0 = df_0.sort_values('Date')

    # 计算: 计算 Cases 的 7 日滚动平均
    df_0['滚动平均(7日)'] = df_0['Cases'].rolling(window=7).mean()

    return df_0
```

[EXPLANATION]

1. **排序**: 按 `Date` 列排序数据。
2. **计算**: 计算 `Cases` 的 7 日滚动平均。
'''

class CodeExplanationAgent(object):

    def __init__(self, client):
        self.client = client

    def run(self, input_tables, code):

        # 生成数据摘要 (包含数据样例)
        data_summary = generate_data_summary(input_tables, include_data_samples=True)

        # 构建用户查询
        user_query = f"[CONTEXT]\n\n{data_summary}\n\n[CODE]\n\nhere is the transformation code: {code}\n\n[EXPLANATION]\n"

        logger.info(user_query)

        # 构建消息队列
        messages = [{"role":"system", "content": SYSTEM_PROMPT},
                    {"role":"user","content": user_query}]

        ###### 调用 OpenAI 接口部分
        response = self.client.get_completion(messages = messages)

        logger.info(f"=== explanation output ===>\n{response.choices[0].message.content}\n")

        return response.choices[0].message.content
